import os
import cv2
import numpy as np

# 文件夹路径
image_dir = "/home/champrin/record_data/网盘/能量机关/2023-联调/标注/5.21/test"       # 文件夹A，存放图片
label_dir = "/home/champrin/record_data/网盘/能量机关/2023-联调/标注/5.21/test"       # 文件夹B，存放对应的YOLO Pose格式标注文件
output_dir = "/home/champrin/record_data/网盘/能量机关/2023-联调/标注/5.21/test/output"      # 可选：变换后保存的位置

os.makedirs(output_dir, exist_ok=True)

# 图像尺寸，YOLO坐标为相对坐标，需要知道原图大小
def yolo_to_pixel(x, y, img_width, img_height):
    return int(float(x) * img_width), int(float(y) * img_height)

# 遍历所有图片
for img_name in os.listdir(image_dir):
    if not img_name.lower().endswith(('.jpg', '.png', '.jpeg')):
        continue

    img_path = os.path.join(image_dir, img_name)
    label_path = os.path.join(label_dir, os.path.splitext(img_name)[0] + ".txt")

    # 判断是否有对应的标注文件
    if not os.path.exists(label_path):
        print(f"跳过 {img_name}，未找到对应标注文件")
        continue

    # 加载图像
    img = cv2.imread(img_path)
    h, w = img.shape[:2]

    with open(label_path, 'r') as f:
        for idx, line in enumerate(f):
            parts = line.strip().split()
            if len(parts) < 14:
                print(f"跳过 {label_path} 的第 {idx+1} 行，关键点不足")
                continue

            # 提取前 4 个关键点坐标（pt1 ~ pt4）
            keypoints = parts[5:13]  # 8 个值，4 对点
            pts = []
            for i in range(0, 8, 2):
                px, py = yolo_to_pixel(keypoints[i], keypoints[i+1], w, h)
                pts.append([px, py])
                cv2.circle(img, (px, py), 2, (0, 255, 0), 1)

            # 相邻两点取中点
            for i in range(len(pts) - 1):
                x1, y1 = pts[i]
                x2, y2 = pts[i + 1]
                mid_x = (x1 + x2) // 2
                mid_y = (y1 + y2) // 2
                cv2.circle(img, (mid_x, mid_y), 2, (0, 0, 255), 1)  # 中点为红色

            # pt1 是第0个点，pt4 是第3个点
            x1, y1 = pts[0]
            x4, y4 = pts[3]
            mid_x = (x1 + x4) // 2
            mid_y = (y1 + y4) // 2

            # 绘制中点（蓝色）
            cv2.circle(img, (mid_x, mid_y), 2, (255, 0, 0), 1)

            src_pts = np.array(pts, dtype=np.float32)

            # 定义目标点坐标（你可以修改为任意矩形）
            dst_pts = np.array([
                [0, 0],
                [200, 0],
                [200, 200],
                [0, 200]
            ], dtype=np.float32)

            # 计算透视变换矩阵
            matrix = cv2.getPerspectiveTransform(src_pts, dst_pts)

            # 应用透视变换
            warped = cv2.warpPerspective(img, matrix, (200, 200))

            # 显示或保存
            out_name = f"{os.path.splitext(img_name)[0]}_warp_{idx+1}.jpg"
            # cv2.imshow(out_name, warped)
            cv2.imshow(out_name, img)
            if cv2.waitKey(0) == 27:
                cv2.destroyAllWindows()
                exit()
            # cv2.imwrite(os.path.join(output_dir, out_name), warped)
            # print(f"已保存: {out_name}")
